// <copyright file="XmlEncryptedTokenSerializer.cs" company="SharpSTS">
// Copyright (c) 2007, 2008 All Right Reserved, http://sharpsts.com/
//
// This source is subject to the Microsoft Permissive License.
// Please see the License.txt file for more information.
// All other rights reserved.
//
// THIS CODE AND INFORMATION ARE PROVIDED "AS IS" WITHOUT WARRANTY OF ANY 
// KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND/OR FITNESS FOR A
// PARTICULAR PURPOSE.
//
// </copyright>
// <author>Barry Dorrans</author>
// <email>barryd@idunno.org</email>
// <date>2008-06-16</date>
// <summary>Encrypts and serializes a token.</summary>

namespace SharpSTS
{
    using System;
    using System.IdentityModel.Selectors;
    using System.IdentityModel.Tokens;
    using System.IO;
    using System.Security.Cryptography;
    using System.Security.Cryptography.X509Certificates;
    using System.ServiceModel.Security;
    using System.ServiceModel.Security.Tokens;
    using System.Xml;

    using Properties;

    /// <summary>
    /// The class used for serializing encrypted SecurityTokens wrapped with a symmetric key,
    /// protected with the public key from an endpoint's X509 certificate.
    /// </summary>
    internal sealed class XmlEncryptedTokenSerializer : WSSecurityTokenSerializer
    {
        #region Defaults
        /// <summary>
        /// The default encryption algorithm to use.
        /// </summary>
        private const string DefaultEncryptionAlgorithm = "http://www.w3.org/2001/04/xmlenc#aes256-cbc";

        /// <summary>
        /// The default key wrapping algorithm to use.
        /// </summary>
        private const string DefaultAsymmetricKeyWrapAlgorithm = "http://www.w3.org/2001/04/xmlenc#rsa-oaep-mgf1p";
        #endregion

        #region Fields
        /// <summary>
        /// The endpoint certificate to encrypt against.
        /// </summary>
        private readonly X509Certificate2 endpointCertificate;

        /// <summary>
        /// The encryption algorithm to use.
        /// </summary>
        private readonly string encryptingAlgorithm;

        /// <summary>
        /// The symmetric encryption process to use.
        /// </summary>
        private readonly SecurityAlgorithmSuite symmetricAlgorithmSuite;

        /// <summary>
        /// The key wrapping algorithm to use.
        /// </summary>
        private readonly string asymmetricKeyWrapAlgorithm;

        /// <summary>
        /// The serializer to use against the token.
        /// </summary>
        private readonly SecurityTokenSerializer tokenSerializer;
        #endregion

        #region Constructors
        /// <summary>
        /// Initializes a new instance of the <see cref="XmlEncryptedTokenSerializer"/> class.
        /// </summary>
        /// <param name="endpointCertificate">The endpoint certificate used to encrypt the token.</param>
        /// <exception cref="ArgumentNullException">Thrown if the EndpointCertificate is null.</exception>
        /// <remarks>If this constructor is used the encryption algorithm defaults to aes256-cbc, and the token serializer defaults to
        /// <see cref="WSSecurityTokenSerializer"/>.</remarks>
        public XmlEncryptedTokenSerializer(X509Certificate2 endpointCertificate) :
            this(endpointCertificate, DefaultEncryptionAlgorithm)
        {            
        }

        /// <summary>
        /// Initializes a new instance of the <see cref="XmlEncryptedTokenSerializer"/> class.
        /// </summary>
        /// <param name="endpointCertificate">The endpoint certificate used to encrypt the token.</param>
        /// <param name="tokenSerializer">The token serializer which produces the unencrypted token.</param>
        /// <exception cref="ArgumentNullException">Thrown if any arguement is null.</exception>
        /// <exception cref="ArgumentException">Thrown if the encrypting algorithm is not known or supported.</exception>
        public XmlEncryptedTokenSerializer(
            X509Certificate2 endpointCertificate,
            SecurityTokenSerializer tokenSerializer) :
            this(endpointCertificate, DefaultEncryptionAlgorithm, DefaultAsymmetricKeyWrapAlgorithm, tokenSerializer)
        {
        }

        /// <summary>
        /// Initializes a new instance of the <see cref="XmlEncryptedTokenSerializer"/> class.
        /// </summary>
        /// <param name="endpointCertificate">The endpoint certificate used to encrypt the token.</param>
        /// <param name="encryptingAlgorithm">The encrypting algorithm to be used; this can be one of
        /// <list type="unordered">
        /// <item>http://www.w3.org/2001/04/xmlenc#aes256-cbc</item>
        /// <item>http://www.w3.org/2001/04/xmlenc#aes192-cbc</item>
        /// <item>http://www.w3.org/2001/04/xmlenc#aes128-cbc</item>
        /// <item>http://www.w3.org/2001/04/xmlenc#tripledes-cbc</item>
        /// </list>
        /// </param>
        /// <exception cref="ArgumentNullException">Thrown if the EndpointCertificate or Encrypting Agorithm is null.</exception>
        /// <exception cref="ArgumentException">Thrown if the encrypting algorithm is not known or supported.</exception>
        /// <remarks>If this constructor is used the key wrap algorithm defaults to rsa-oaep-mgf1p.</remarks>
        public XmlEncryptedTokenSerializer(
            X509Certificate2 endpointCertificate, 
            string encryptingAlgorithm)
            : this(endpointCertificate, encryptingAlgorithm, DefaultAsymmetricKeyWrapAlgorithm, new WSSecurityTokenSerializer())
        {       
        }

        /// <summary>
        /// Initializes a new instance of the <see cref="XmlEncryptedTokenSerializer"/> class.
        /// </summary>
        /// <param name="endpointCertificate">The endpoint certificate used to encrypt the token.</param>
        /// <param name="encryptingAlgorithm">The encrypting algorithm to be used; this can be one of
        /// <list type="unordered">
        /// <item>http://www.w3.org/2001/04/xmlenc#aes256-cbc</item>
        /// <item>http://www.w3.org/2001/04/xmlenc#aes192-cbc</item>
        /// <item>http://www.w3.org/2001/04/xmlenc#aes128-cbc</item>
        /// <item>http://www.w3.org/2001/04/xmlenc#tripledes-cbc</item>
        /// </list>
        /// </param>
        /// <param name="asymmetricKeyWrapAlgorithm">The asymmetric key wrap algorithm to be used.</param>
        /// <exception cref="ArgumentNullException">Thrown if the EndpointCertificate or Encrypting Agorithm is null.</exception>
        /// <exception cref="ArgumentException">Thrown if the encrypting algorithm is not known or supported.</exception>
        /// <remarks>If this constructor is used the key wrap algorithm defaults to rsa-oaep-mgf1p.</remarks>
        public XmlEncryptedTokenSerializer(X509Certificate2 endpointCertificate, string encryptingAlgorithm, string asymmetricKeyWrapAlgorithm)
            : this(endpointCertificate, encryptingAlgorithm, asymmetricKeyWrapAlgorithm, new WSSecurityTokenSerializer())
        {
        }

        /// <summary>
        /// Initializes a new instance of the <see cref="XmlEncryptedTokenSerializer"/> class.
        /// </summary>
        /// <param name="endpointCertificate">The endpoint certificate used to encrypt the token.</param>
        /// <param name="encryptingAlgorithm">The encrypting algorithm to be used; this can be one of
        /// <list type="unordered">
        /// <item>http://www.w3.org/2001/04/xmlenc#aes256-cbc</item>
        /// <item>http://www.w3.org/2001/04/xmlenc#aes192-cbc</item>
        /// <item>http://www.w3.org/2001/04/xmlenc#aes128-cbc</item>
        /// <item>http://www.w3.org/2001/04/xmlenc#tripledes-cbc</item>
        /// </list>
        /// </param>
        /// <param name="asymmetricKeyWrapAlgorithm">The asymmetric key wrap algorithm to be used.</param>
        /// <param name="tokenSerializer">The token serializer which produces the unencrypted token.</param>
        /// <exception cref="ArgumentNullException">Thrown any argument is null.</exception>
        /// <exception cref="ArgumentException">Thrown if the encrypting algorithm is not known or supported.</exception>
        public XmlEncryptedTokenSerializer(
            X509Certificate2 endpointCertificate, 
            string encryptingAlgorithm, 
            string asymmetricKeyWrapAlgorithm,
            SecurityTokenSerializer tokenSerializer)
        {
            if (null == endpointCertificate)
            {
                throw new ArgumentNullException("endpointCertificate");
            }

            this.endpointCertificate = endpointCertificate;

            if (string.IsNullOrEmpty(encryptingAlgorithm))
            {
                throw new ArgumentNullException("encryptingAlgorithm");
            }
            
            this.encryptingAlgorithm = encryptingAlgorithm;

            if (string.IsNullOrEmpty(asymmetricKeyWrapAlgorithm))
            {
                throw new ArgumentNullException("asymmetricKeyWrapAlgorithm");
            }
            
            this.asymmetricKeyWrapAlgorithm = asymmetricKeyWrapAlgorithm;

            if (string.IsNullOrEmpty(encryptingAlgorithm))
            {
                throw new ArgumentNullException("tokenSerializer");
            }
            
            this.tokenSerializer = tokenSerializer;

            switch (encryptingAlgorithm)
            {
                case "http://www.w3.org/2001/04/xmlenc#aes256-cbc":
                    this.symmetricAlgorithmSuite = SecurityAlgorithmSuite.Basic256;
                    break;
                case "http://www.w3.org/2001/04/xmlenc#aes128-cbc":
                    this.symmetricAlgorithmSuite = SecurityAlgorithmSuite.Basic128;
                    break;
                case "http://www.w3.org/2001/04/xmlenc#aes192-cbc":
                    this.symmetricAlgorithmSuite = SecurityAlgorithmSuite.Basic192;
                    break;
                case "http://www.w3.org/2001/04/xmlenc#tripledes-cbc":
                    this.symmetricAlgorithmSuite = SecurityAlgorithmSuite.TripleDes;
                    break;
                default:
                    throw new ArgumentException(Resources.ErrorUnknownEncryption, "encryptingAlgorithm");
            }
        }
        #endregion

        #region Serializer methods
        /// <summary>
        /// Writes the token core with a specified security token using the specified <see cref="T:System.Xml.XmlWriter"/>.
        /// </summary>
        /// <param name="writer">The specified <see cref="T:System.Xml.XmlWriter"/>.</param>
        /// <param name="token">The <see cref="T:System.IdentityModel.Tokens.SecurityToken"/>.</param>
        /// <exception cref="T:System.Xml.XmlException">An error occurred when serializing the security token.</exception>
        /// <exception cref="T:System.InvalidOperationException">The token was not successfully written.</exception>
        protected override void WriteTokenCore(XmlWriter writer, SecurityToken token)
        {
            SecurityTokenSerializer baseSerializer = this.tokenSerializer;
            XmlElement encryptedToken;

            using (MemoryStream stream = new MemoryStream())
            {
                XmlDictionaryWriter baseWriter =
                    XmlDictionaryWriter.CreateDictionaryWriter(new XmlTextWriter(new StreamWriter(stream)));
                baseSerializer.WriteToken(baseWriter, token);
                baseWriter.Flush();
                stream.Seek(0L, SeekOrigin.Begin);
                encryptedToken = this.EncryptToken(stream);
                baseWriter.Close();
                stream.Close();
            }

            encryptedToken.WriteTo(writer);
        }
        #endregion

        #region Encryption
        /// <summary>
        /// Creates an <see cref="XmlDictionaryReader"/> with suitable default quotas.
        /// </summary>
        /// <param name="input">The stream to create the reader from.</param>
        /// <returns>A quota bound <see cref="XmlDictionaryReader"/> from the specified stream.</returns>
        private static XmlDictionaryReader CreateReaderWithQuotas(Stream input)
        {
            XmlDictionaryReaderQuotas defaultQuotas = new XmlDictionaryReaderQuotas
            {
                MaxDepth = 0x20,
                MaxStringContentLength = 0x2000,
                MaxArrayLength = 0x1400000,
                MaxBytesPerRead = 0x1000,
                MaxNameTableCharCount = 0x4000
            };

            return XmlDictionaryReader.CreateTextReader(input, null, defaultQuotas, null);
        }

        /// <summary>
        /// Encrypts the specified memory stream in an Xml encryption wrapper.
        /// </summary>
        /// <param name="streamToBeEncrypted">The memory stream to be encrypted.</param>
        /// <returns>An XmlElement containing the encrypted memory stream.</returns>
        private XmlElement EncryptToken(MemoryStream streamToBeEncrypted)
        {
            SecurityToken wrappingToken = new X509SecurityToken(this.endpointCertificate, "id");
            SecurityKeyIdentifier wrappingTokenReference =
                new SecurityKeyIdentifier(new SecurityKeyIdentifierClause[]
                                              {
                                                  wrappingToken.CreateKeyIdentifierClause<X509ThumbprintKeyIdentifierClause>()
                                              });
            int num = this.symmetricAlgorithmSuite.DefaultEncryptionKeyDerivationLength / 8;
            byte[] buffer = new byte[num];
            new RNGCryptoServiceProvider().GetNonZeroBytes(buffer);
            WrappedKeySecurityToken token = new WrappedKeySecurityToken(
                string.Empty, 
                buffer,
                this.asymmetricKeyWrapAlgorithm, 
                wrappingToken, 
                wrappingTokenReference);
            SecurityKeyIdentifier identifier = new SecurityKeyIdentifier(new SecurityKeyIdentifierClause[]
                                                                             {
                                                                                 new EncryptedKeyIdentifierClause(
                                                                                     token.GetWrappedKey(), 
                                                                                     token.WrappingAlgorithm, 
                                                                                     token.WrappingTokenReference)
                                                                             });
            SymmetricAlgorithm symmetricAlgorithm = ((SymmetricSecurityKey)token.SecurityKeys[0]).GetSymmetricAlgorithm(this.encryptingAlgorithm);
            EncryptedData data = new EncryptedData
            {
                TokenSerializer = new WSSecurityTokenSerializer(true),
                KeyIdentifier = identifier,
                Algorithm = this.encryptingAlgorithm,
                Type = "http://www.w3.org/2001/04/xmlenc#Element"
            };

            data.SetUpEncryption(symmetricAlgorithm, streamToBeEncrypted.GetBuffer(), 0, Convert.ToInt32(streamToBeEncrypted.Length));

            XmlElement element;
            using (MemoryStream stream = new MemoryStream())
            {
                XmlDictionaryWriter writer =
                    XmlDictionaryWriter.CreateDictionaryWriter(new XmlTextWriter(new StreamWriter(stream)));
                data.WriteTo(writer);
                writer.Flush();
                stream.Seek(0L, SeekOrigin.Begin);
                XmlDocument document = new XmlDocument();
                element = (XmlElement)document.ReadNode(CreateReaderWithQuotas(stream));
                Array.Clear(streamToBeEncrypted.GetBuffer(), 0, Convert.ToInt32(streamToBeEncrypted.Length));
                Array.Clear(stream.GetBuffer(), 0, Convert.ToInt32(stream.Length));
                streamToBeEncrypted.Close();
                stream.Close();
            }

            return element;
        }
        #endregion
    }
}
